/*
Copyright (c) 2007, Miron Brezuleanu

All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
* Neither the name of the product nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
*/


// Mon.cpp : Defines the entry point for the application.
//

#include "stdafx.h"
#include <pdh.h>
#include "wmon.h"
#include "estrbuf.h"
#include "PdhWrapper.h"
#include <vector>
#include <shlwapi.h>
#include <time.h>
#include "monitor.h"
#include <map>

#pragma comment(lib, "shlwapi")

using namespace std;

#define MAX_LOADSTRING 100

// Global Variables:
HINSTANCE hInst;								// current instance
TCHAR szTitle[MAX_LOADSTRING];					// The title bar text
TCHAR szWindowClass[MAX_LOADSTRING];			// the main window class name

// Forward declarations of functions included in this code module:
ATOM				MyRegisterClass(HINSTANCE hInstance);
BOOL				InitInstance(HINSTANCE, int);
LRESULT CALLBACK	WndProc(HWND, UINT, WPARAM, LPARAM);
INT_PTR CALLBACK	AboutDlgProc(HWND, UINT, WPARAM, LPARAM);
INT_PTR CALLBACK	ConfigDlgProc(HWND, UINT, WPARAM, LPARAM);

PdhWrapper pdhWrapper;

class WMonConfig
{
	static const DWORD magic = 0xC078AEED;
	static const DWORD version = 0x00000001;

	static void GetConfigFileName(tstring &fileName)
	{
		TCHAR path[MAX_PATH];
		GetModuleFileName(NULL, path, MAX_PATH);
		int len = (int)_tcslen(path) - 1;
		while (len >= 0 && path[len] != _T('\\'))
		{
			len --;
		}
		if (len >= 0)
		{
			path[len] = _T('\0');
			fileName = path;
			fileName.append(_T("\\wmon.cfg"));
		}
		else
		{
			fileName = _T("c:\\wmon.cfg");
		}
	}

	int m_frequency;
	int m_winX, m_winY;
	time_t m_lastSaveTime;
	bool m_dirty;
	tstring m_fileName;

public:
	void SetFrequency(int frequency)
	{
		if (m_frequency != frequency) {
			m_frequency = frequency;
			m_dirty = true;
		}
	}

	int GetFrequency()
	{
		return this->m_frequency;
	}

	void SetWindowPosition(int x, int y)
	{
		if (m_winX != x || m_winY != y)
		{
			m_winX = x;
			m_winY = y;
			m_dirty = true;
		}
	}

	void GetWindowPosition(int &x, int &y)
	{
		x = m_winX;
		y = m_winY;
	}

	WMonConfig()
	{
		GetConfigFileName(m_fileName);
		bool readOk = false;
		if (PathFileExists(m_fileName.c_str())) {
			FILE *fp = _tfopen(m_fileName.c_str(), _T("rb"));
			if (fp)
			{
				DWORD ver, mag;
				fread(&mag, sizeof(mag), 1, fp);
				if (mag == WMonConfig::magic)
				{
					fread(&ver, sizeof(ver), 1, fp);
					if (ver == WMonConfig::version)
					{
						fread(&m_frequency, sizeof(m_frequency), 1, fp);
						fread(&m_winX, sizeof(m_winX), 1, fp);
						fread(&m_winY, sizeof(m_winY), 1, fp);						
						readOk = true;
					}
				}
				fclose(fp);
			}
		}
		if (!readOk)
		{
			m_frequency = 10;
			m_winX = 100;
			m_winY = 100;
			m_dirty = true;
		}
		m_lastSaveTime = time(NULL);
	}

	~WMonConfig()
	{
		Save();
	}

	void Save()
	{
		time_t currentTime = time(NULL);
		if (m_dirty && currentTime > m_lastSaveTime + 1)
		{
			FILE *fp = _tfopen(m_fileName.c_str(), _T("wb"));
			if (fp)
			{
				fwrite(&magic, sizeof(WMonConfig::magic), 1, fp);
				fwrite(&version, sizeof(WMonConfig::version), 1, fp);
				fwrite(&m_frequency, sizeof(m_frequency), 1, fp);
				fwrite(&m_winX, sizeof(m_winX), 1, fp);
				fwrite(&m_winY, sizeof(m_winY), 1, fp);
				m_dirty = false;
				fclose(fp);
			}
			m_lastSaveTime = currentTime;
		}
	}
};

const int monitorHeight = 40;
const int monitorWidth = 60;

WMonConfig config;
HWND mainWindowHandle = NULL;
map<tstring, Monitor *> monitors;

void
ConfigurePdhWrapper()
{
	SYSTEM_INFO si;
	GetSystemInfo(&si);
	for (unsigned i = 0; i < si.dwNumberOfProcessors; i++)
	{
		DWORD cpuMask = 1 << i;
		if (si.dwActiveProcessorMask & cpuMask)
		{
			EStrBuf name(100), pdhPath(100);
			name.extendf(_T("cpu%d"), i);
			pdhPath.extendf(_T("\\Processor(%d)\\%% Processor Time"), i);
			pdhWrapper.AddCounter(name.c_str(), pdhPath.c_str());
			monitors[tstring(name.c_str())] = new AveragingPercentMonitor(5, monitorWidth, monitorHeight, name.c_str());
		}
	}
	pdhWrapper.AddCounter(_T("mem"), _T("\\Memory\\% Committed Bytes In Use"));
	monitors[tstring(_T("mem"))] = new AveragingPercentMonitor(5, monitorWidth, monitorHeight, _T("mem"));
}

void
FitInWorkarea(HWND hWnd)
{
	RECT workareaRect;
	SystemParametersInfo(SPI_GETWORKAREA, 0, &workareaRect, 0);
	RECT windowRect;
	GetWindowRect(hWnd, &windowRect);
	int winWidth, winHeight;
	winWidth = windowRect.right - windowRect.left;
	winHeight = windowRect.bottom - windowRect.top;
	if (windowRect.left < workareaRect.left)
	{
		windowRect.left = workareaRect.left;
	}
	if (windowRect.top < workareaRect.top)
	{
		windowRect.top = workareaRect.top;
	}
	if (windowRect.left + winWidth > workareaRect.right)
	{
		windowRect.left = workareaRect.right - winWidth;
	}
	if (windowRect.top + winHeight > workareaRect.bottom)
	{
		windowRect.top = workareaRect.bottom - winHeight;
	}
	RECT oldWindowRect;
	GetWindowRect(hWnd, &oldWindowRect);
	if (oldWindowRect.left != windowRect.left || oldWindowRect.top != windowRect.top)
	{
		SetWindowPos(hWnd, HWND_TOPMOST, windowRect.left, windowRect.top, winWidth, winHeight, 0);
		config.SetWindowPosition(windowRect.left, windowRect.top);
	}
}

int APIENTRY _tWinMain(HINSTANCE hInstance,
					   HINSTANCE hPrevInstance,
					   LPTSTR    lpCmdLine,
					   int       nCmdShow)
{
	ConfigurePdhWrapper();

	UNREFERENCED_PARAMETER(hPrevInstance);
	UNREFERENCED_PARAMETER(lpCmdLine);

	MSG msg;
	HACCEL hAccelTable;

	// Initialize global strings
	LoadString(hInstance, IDS_APP_TITLE, szTitle, MAX_LOADSTRING);
	LoadString(hInstance, IDC_WMON, szWindowClass, MAX_LOADSTRING);
	MyRegisterClass(hInstance);

	// Perform application initialization:
	if (!InitInstance (hInstance, nCmdShow))
	{
		return FALSE;
	}

	hAccelTable = LoadAccelerators(hInstance, MAKEINTRESOURCE(IDC_WMON));

	// Main message loop:
	while (GetMessage(&msg, NULL, 0, 0))
	{
		if (!TranslateAccelerator(msg.hwnd, hAccelTable, &msg))
		{
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
	}

	return (int) msg.wParam;
}

//
//  FUNCTION: MyRegisterClass()
//
//  PURPOSE: Registers the window class.
//
//  COMMENTS:
//
//    This function and its usage are only necessary if you want this code
//    to be compatible with Win32 systems prior to the 'RegisterClassEx'
//    function that was added to Windows 95. It is important to call this function
//    so that the application will get 'well formed' small icons associated
//    with it.
//
ATOM MyRegisterClass(HINSTANCE hInstance)
{
	WNDCLASSEX wcex;

	wcex.cbSize = sizeof(WNDCLASSEX);

	wcex.style			= CS_HREDRAW | CS_VREDRAW;
	wcex.lpfnWndProc	= WndProc;
	wcex.cbClsExtra		= 0;
	wcex.cbWndExtra		= 0;
	wcex.hInstance		= hInstance;
	wcex.hIcon			= LoadIcon(hInstance, MAKEINTRESOURCE(IDI_WMON));
	wcex.hCursor		= LoadCursor(NULL, IDC_ARROW);
	wcex.hbrBackground	= (HBRUSH)(COLOR_WINDOW+1);
	wcex.lpszMenuName	= NULL; //MAKEINTRESOURCE(IDC_MON);
	wcex.lpszClassName	= szWindowClass;
	wcex.hIconSm		= LoadIcon(wcex.hInstance, MAKEINTRESOURCE(IDI_SMALL));

	return RegisterClassEx(&wcex);
}

void
ResetTimer()
{
	const int timerId = 1;
	static bool hasTimer = false;
	if (mainWindowHandle)
	{
		if (hasTimer) 
		{
			KillTimer(mainWindowHandle, timerId);
		}
		SetTimer(mainWindowHandle, timerId, 1000 / config.GetFrequency(), NULL);
		hasTimer = true;
	}
}

//
//   FUNCTION: InitInstance(HINSTANCE, int)
//
//   PURPOSE: Saves instance handle and creates main window
//
//   COMMENTS:
//
//        In this function, we save the instance handle in a global variable and
//        create and display the main program window.
//
BOOL InitInstance(HINSTANCE hInstance, int nCmdShow)
{
	HWND hWnd;

	hInst = hInstance; // Store instance handle in our global variable

	vector<tstring> pdhCounters;
	pdhWrapper.EnumerateCounters(pdhCounters);
	int winX, winY;
	config.GetWindowPosition(winX, winY);
	HWND invisibleParent = CreateWindowEx(WS_EX_APPWINDOW, // change this to WS_EX_TOOLWINDOW to hide from alt tab
		szWindowClass,
		_T(""),
		WS_POPUP,
		CW_USEDEFAULT,
		CW_USEDEFAULT,
		CW_USEDEFAULT,
		CW_USEDEFAULT,
		NULL,
		NULL,
		hInstance,
		0);
	hWnd = CreateWindow(szWindowClass, 
		szTitle, 
		WS_POPUP | WS_BORDER, //WS_OVERLAPPEDWINDOW,
		winX, 
		winY, 
		monitorWidth, 
		monitorHeight * (int)pdhCounters.size(), 
		invisibleParent, 
		NULL, 
		hInstance, 
		NULL);
	mainWindowHandle = hWnd;

	if (!hWnd)
	{
		return FALSE;
	}

	ResetTimer();
	ShowWindow(hWnd, nCmdShow);
	UpdateWindow(hWnd);
	SetWindowPos(hWnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE);

	return TRUE;
}

void
DrawError(HWND hWnd, HDC hdc)
{
	HPEN hpen = CreatePen(PS_SOLID, 1, RGB(0xff, 0x00, 0x00));
	HGDIOBJ oldPen = SelectObject(hdc, hpen);
	RECT clientRect;
	GetClientRect(hWnd, &clientRect);
	MoveToEx(hdc, clientRect.left, clientRect.top, NULL);
	LineTo(hdc, clientRect.right, clientRect.bottom);
	MoveToEx(hdc, clientRect.right - 1, clientRect.top, NULL);
	LineTo(hdc, clientRect.left, clientRect.bottom);
	SelectObject(hdc, oldPen);
	DeleteObject(hpen);
}

void
PaintStuffImpl(HWND hWnd, HDC hdc)
{
	RECT rect;
	GetClientRect(hWnd, &rect);
	LOGBRUSH cfgBgBrush;
	cfgBgBrush.lbColor = RGB(0xFF, 0xFF, 0x80);
	cfgBgBrush.lbHatch = 0;
	cfgBgBrush.lbStyle = BS_SOLID;
	HBRUSH bgBrush = CreateBrushIndirect(&cfgBgBrush);
	FillRect(hdc, &rect, bgBrush);
	DeleteObject(bgBrush);
	SelectObject(hdc, GetStockObject(DEFAULT_GUI_FONT));
	GetClientRect(hWnd, &rect);
	pdhWrapper.ReadCounters();
	if (!pdhWrapper.IsOk()) {
		pdhWrapper.Reset();
		pdhWrapper.ReadCounters();
		if (!pdhWrapper.IsOk()) {
			DrawError(hWnd, hdc);
			return;
		}
	}
	vector<tstring> nameList;
	pdhWrapper.EnumerateCounters(nameList);
	int yOffset = 0;
	SetBkMode(hdc, TRANSPARENT);
	for (unsigned i = 0; i < nameList.size(); ++i)
	{
		PDH_FMT_COUNTERVALUE counterValue;
		pdhWrapper.GetCounterValue(nameList[i], counterValue);
		if (pdhWrapper.IsOk())
		{
			Monitor *mon = monitors[nameList[i]];
			mon->AddSample(counterValue.longValue / 100.0);
			mon->Display(hdc, 0, yOffset);
			RECT bbox;
			mon->GetBoundingBox(bbox);
			yOffset += bbox.bottom - bbox.top;
		}
	}
}

void
PaintStuff(HWND hWnd, HDC hdc)
{
	RECT rect;
	GetClientRect(hWnd, &rect);
	HBITMAP bs = CreateCompatibleBitmap(hdc, rect.right - rect.left, rect.bottom - rect.top);
	HDC hbsdc = CreateCompatibleDC(hdc);
	HGDIOBJ oldMemBmp = SelectObject(hbsdc, bs);
	PaintStuffImpl(hWnd, hbsdc);
	BitBlt(hdc, 0, 0, rect.right - rect.left, rect.bottom - rect.top, hbsdc, 0, 0, SRCCOPY);
	SelectObject(hbsdc, oldMemBmp);
	DeleteDC(hbsdc);
	DeleteObject(bs);
}

void
UpdateDisplay(HWND hWnd)
{
	HDC dc = GetDC(hWnd);
	PaintStuff(hWnd, dc);
	ReleaseDC(hWnd, dc);
}

void
GetDiffPt(POINT *pt, int client_x, int client_y, HWND hWnd)
{
	POINT screenpt;
	screenpt.x = client_x;
	screenpt.y = client_y;
	ClientToScreen(hWnd, &screenpt);
	RECT wRect;
	GetWindowRect(hWnd, &wRect);
	pt->x = wRect.left - screenpt.x;
	pt->y = wRect.top - screenpt.y;
}

void
WMonMoveWindow(HWND hWnd, POINT *diff_pt, int x, int y)
{
	POINT screenpt;
	screenpt.x = x;
	screenpt.y = y;
	ClientToScreen(hWnd, &screenpt);
	POINT winpos;
	winpos.x = screenpt.x + diff_pt->x;
	winpos.y = screenpt.y + diff_pt->y;	
	bool adjust = false;
	RECT warect;
	SystemParametersInfo(SPI_GETWORKAREA, 0, &warect, 0);
	RECT winrect;
	GetWindowRect(hWnd, &winrect);
	if (winpos.x < warect.left)
	{
		winpos.x = warect.left;
		adjust = true;
	}
	if (winpos.y < warect.top)
	{
		winpos.y = warect.top;
		adjust = true;
	}
	if (winpos.x + winrect.right - winrect.left > warect.right)
	{
		winpos.x = warect.right + winrect.left - winrect.right;
		adjust = true;
	}
	if (winpos.y + winrect.bottom - winrect.top > warect.bottom)
	{
		winpos.y = warect.bottom + winrect.top - winrect.bottom;
		adjust = true;
	}
	// LOWORD/HIWORD kill the sign of the coordinate, need to take care of it here...
	if (winpos.x > 30000 || winpos.y > 30000) 
	{
		adjust = true;
	}
	else
	{
		SetWindowPos(hWnd, NULL, winpos.x, winpos.y, 0, 0, SWP_NOZORDER | SWP_NOSIZE);
		config.SetWindowPosition(winpos.x, winpos.y);
	}
	if (adjust)
	{
		GetDiffPt(diff_pt, x, y, hWnd);
	}
}

void
ShowContextMenu(HWND hWnd, int x, int y)
{
	const int menuAbout = 1;
	const int menuExit = 2;
	const int menuConfig = 3;
	POINT screenpt;
	screenpt.x = x;
	screenpt.y = y;
	ClientToScreen(hWnd, &screenpt);
	HMENU hMenu = CreatePopupMenu();
	AppendMenu(hMenu, MF_STRING, menuConfig, _T("Config"));
	AppendMenu(hMenu, MF_SEPARATOR, 0, NULL);
	AppendMenu(hMenu, MF_STRING, menuAbout, _T("About"));
	AppendMenu(hMenu, MF_SEPARATOR, 0, NULL);
	AppendMenu(hMenu, MF_STRING, menuExit, _T("Exit"));
	int commandId = TrackPopupMenu(hMenu, 
		TPM_RIGHTBUTTON | TPM_RETURNCMD, 
		screenpt.x, 
		screenpt.y, 
		0, 
		hWnd, 
		NULL);
	DestroyMenu(hMenu);
	switch (commandId)
	{
	case menuAbout:
		DialogBox(hInst, MAKEINTRESOURCE(IDD_ABOUTBOX), hWnd, AboutDlgProc);
		break;
	case menuConfig:
		DialogBox(hInst, MAKEINTRESOURCE(IDD_CONFIG), hWnd, ConfigDlgProc);
		break;
	case menuExit:
		PostQuitMessage(0);
		break;
	}
}

//
//  FUNCTION: WndProc(HWND, UINT, WPARAM, LPARAM)
//
//  PURPOSE:  Processes messages for the main window.
//
//  WM_COMMAND	- process the application menu
//  WM_PAINT	- Paint the main window
//  WM_DESTROY	- post a quit message and return
//
//
LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
	static bool moving = false;
	static POINT diff_pt;
	PAINTSTRUCT ps;
	HDC hdc;

	switch (message)
	{
	case WM_PAINT:
		hdc = BeginPaint(hWnd, &ps);
		PaintStuff(hWnd, hdc);
		EndPaint(hWnd, &ps);
		break;
	case WM_DESTROY:
		PostQuitMessage(0);
		break;
	case WM_LBUTTONDOWN:
		moving = true;
		GetDiffPt(&diff_pt, LOWORD(lParam), HIWORD(lParam), hWnd);
		SetCapture(hWnd);
		break;
	case WM_MOUSEMOVE:
		if (moving) 
		{
			WMonMoveWindow(hWnd, &diff_pt, LOWORD(lParam), HIWORD(lParam));
		}
		break;
	case WM_LBUTTONUP:
		moving = false;
		ReleaseCapture();
		break;
	case WM_RBUTTONUP:
		ShowContextMenu(hWnd, LOWORD(lParam), HIWORD(lParam));
		break;
	case WM_TIMER:
		FitInWorkarea(hWnd);
		UpdateDisplay(hWnd);
		config.Save();
		break;
	default:
		return DefWindowProc(hWnd, message, wParam, lParam);
	}
	return 0;
}

// Message handler for about box.
INT_PTR CALLBACK AboutDlgProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	UNREFERENCED_PARAMETER(lParam);
	switch (message)
	{
	case WM_INITDIALOG:
		return (INT_PTR)TRUE;

	case WM_COMMAND:
		if (LOWORD(wParam) == IDOK || LOWORD(wParam) == IDCANCEL)
		{
			EndDialog(hDlg, LOWORD(wParam));
			return (INT_PTR)TRUE;
		}
		break;
	}
	return (INT_PTR)FALSE;
}

// Message handler for config box.
INT_PTR CALLBACK ConfigDlgProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	UNREFERENCED_PARAMETER(lParam);
	DWORD cmd;
	switch (message)
	{
	case WM_INITDIALOG:
		{
			EStrBuf buf(100);
			buf.extendf(_T("%ld"), config.GetFrequency());
			HWND tbHandle = GetDlgItem(hDlg, IDC_EDITFREQUENCY);
			if (tbHandle)
			{
				SetWindowText(tbHandle, buf.c_str());
			}
		}
		return (INT_PTR)TRUE;

	case WM_COMMAND:
		{
			cmd = LOWORD(wParam);
			if (cmd == IDOK || cmd == IDCANCEL)
			{
				bool exitOk = true;
				if (cmd == IDOK) 
				{
					const int maxChars = 100;
					TCHAR buf[maxChars];
					HWND tbHandle = GetDlgItem(hDlg, IDC_EDITFREQUENCY);
					if (tbHandle)
					{
						long value;
						GetWindowText(tbHandle, buf, maxChars);
						_stscanf_s(buf, _T("%ld"), &value);
						if (value < 1 || value > 10)
						{
							HWND erStatic = GetDlgItem(hDlg, IDC_STATICERROR);
							if (erStatic) {
								SetWindowText(erStatic,
									_T("Must enter a value between 1 and 10"));
							}
							exitOk = false;
						}
						else
						{
							config.SetFrequency(value);
							ResetTimer(); 
						}
					}
				}
				if (exitOk)
				{
					EndDialog(hDlg, LOWORD(wParam));
				}
				return (INT_PTR)TRUE;
			}
		}
		break;
	}
	return (INT_PTR)FALSE;
}
